import os
import pandas as pd
from PIL import Image
from paddle.io import Dataset
from paddle.vision.transforms import ToTensor
from utils.process_labels import encode_labels

# crop the image to discard useless parts
def crop_data(image, label, offset=690):
    """
    Attention:
    h,w,c = image.shape
    """
    height = image.height
    width = image.width
    return image.crop((0, offset, width, height)), label.crop((0, offset, width, height))

class LaneDataset(Dataset):

    def __init__(self, csv_file, transform=None):
        super(LaneDataset, self).__init__()
        self.data = pd.read_csv(os.path.join(os.getcwd(), "work", "data_list", csv_file), header=None, sep=' ',
                                  names=["image","label"])
        self.images = self.data["image"].values[1:]
        self.labels = self.data["label"].values[1:]

        self.transform = transform

    def __len__(self):
        return self.images.shape[0]

    def __getitem__(self, idx):
        train_img = Image.open(self.images[idx])
        train_mask = Image.open(self.labels[idx])
        train_img, train_mask = crop_data(train_img, train_mask)
        # Encode
        train_mask = encode_labels(train_mask)
        sample = [train_img, train_mask]
        if self.transform:
            train_img = self.transform(train_img)
            train_mask = self.transform(train_mask)
            sample = [train_img, train_mask]
        return sample

if __name__ == '__main__':
    # custom_dataset = LaneDataset(csv_file, transform=[ToTensor()])
    train_img = Image.open(r"D:\nlpStudy\cv_study\project1_data\ColorImage\ColorImage_road02\ColorImage\Record001\Camera_5\170927_063811892_Camera_5.jpg")
    height = train_img.height
    width = train_img.width
    train_img = train_img.crop((0, 690, width, height))
    train_img.save(r"D:\nlpStudy\cv_study\project1_data\ColorImage\ColorImage_road02\ColorImage\Record001\Camera_5\170927_063811892_Camera_5_cut.jpg")
